"""
多模态调用
"""
from marker.schema import BlockTypes
from typing import List
from marker.schema.blocks import Block
from PIL import Image
from tqdm import tqdm
from concurrent.futures import ThreadPoolExecutor, as_completed

from ..processors import BaseOptProcessor
from ..model.qwen_vl import QwenVL

class OptImageDescriptionProcessor(BaseOptProcessor):
    block_types = (
        BlockTypes.Picture,
        BlockTypes.Figure,
    )
    qwen = QwenVL()

    def process_rewriting(self,
            doc_blocks: List[Block],
            block_images: List[Image.Image]
    ):
        pbar = tqdm(
            desc=f"{self.__class__.__name__} running", disable=self.disable_tqdm
        )
        with ThreadPoolExecutor(max_workers=self.max_concurrency) as executor:
            for future in as_completed(
                [
                    executor.submit(self.single_process, block, image)
                    for block, image in zip(doc_blocks, block_images)
                ]
            ):
                future.result()
                pbar.update(1)
        pbar.close()

    def single_process(self, block: Block, image: Image.Image):
        vl_res = self.qwen.invoke(image)
        block.html = f"<p>图片描述：{vl_res['text']}</p>"